
#include "replication_pad2d_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

constexpr uint32_t BLOCK_SIZE = 32;

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    ReplicationPad2dTilingData tiling;
    const gert::StorageShape* x_shape = context->GetInputShape(0);
    int32_t shapedimnum = x_shape->GetStorageShape().GetDimNum();
    int32_t lastdimlength = x_shape->GetStorageShape().GetDim(shapedimnum - 1);
    int32_t lastB2dimlength = x_shape->GetStorageShape().GetDim(shapedimnum - 2);
    int32_t totallength = 1;
    for (int i = 0; i < x_shape->GetStorageShape().GetDimNum(); i++)
        totallength *= x_shape->GetStorageShape().GetDim(i);
    tiling.set_lastdimlength(lastdimlength);
    tiling.set_lastB2dimlength(lastB2dimlength);
    tiling.set_totallength(totallength);

    uint64_t ubSize;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
    uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize(); 
    uint32_t inputBytes = GetSizeByDataType(context->GetInputDesc(0)->GetDataType()); 
    uint32_t inputLength = inputBytes * inputNum;
    uint32_t ubDataNumber = (inputBytes == 2) ? 12 : 12;
    uint32_t tileBlockNum = (ubSize / BLOCK_SIZE) / ubDataNumber; //每个ub段可用的空间块数
    uint32_t tileDataNum = (tileBlockNum * BLOCK_SIZE) / inputBytes; //每次处理的数据量
    tiling.set_tileDataNum(tileDataNum);
    context->SetTilingKey(100);

    const gert::Tensor *pad_tensor = context->GetInputTensor(1);
    if (x_shape == nullptr || pad_tensor == nullptr) {
        return ge::GRAPH_FAILED;
    }
    auto addrL = pad_tensor->GetData<int32_t>();
    int32_t leftpad = *addrL;
    auto addrR = pad_tensor->GetData<int32_t>() + 1;
    int32_t rightpad = *addrR;    
    auto addrT = pad_tensor->GetData<int32_t>() + 2;
    int32_t toppad = *addrT;  
    auto addrB = pad_tensor->GetData<int32_t>() + 3;
    int32_t bottompad = *addrB;  
    tiling.set_leftpad(leftpad);
    tiling.set_rightpad(rightpad);
    tiling.set_toppad(toppad);
    tiling.set_bottompad(bottompad);

    int32_t highdimcount = totallength / lastdimlength / lastB2dimlength;
    int32_t outputlength = (lastdimlength + leftpad + rightpad) * (lastB2dimlength + toppad + bottompad) * highdimcount;
    tiling.set_outputlength(outputlength);
    tiling.set_highdimcount(highdimcount);

    int32_t outputline = lastdimlength + leftpad + rightpad;
    int32_t alignnum = 32 / inputBytes;
    int32_t outputlinealign = (outputline + alignnum - 1) / alignnum * alignnum;
    tiling.set_outputlinealign(outputlinealign);
    if(lastdimlength > tileDataNum) {
        context->SetTilingKey(200);
    } else{
        if((lastdimlength * inputBytes % 32 == 0) && (leftpad * inputBytes % 32 == 0)) {
            context->SetTilingKey(105);
            ubDataNumber = (inputBytes == 2) ? 4 : 4;
            tileBlockNum = (ubSize / BLOCK_SIZE) / ubDataNumber; //每个ub段可用的空间块数
            tileDataNum = (tileBlockNum * BLOCK_SIZE) / inputBytes; //每次处理的数据量
            tiling.set_tileDataNum(tileDataNum);
            if(rightpad * inputBytes % 32 == 0) {
                context->SetTilingKey(106);
            }
        }
    }

    std::cout << "tiling_key:" << context->GetTilingKey() << std::endl;
    std::cout << "lastdimlength:" << lastdimlength << std::endl;
    std::cout << "lastB2dimlength:" << lastB2dimlength << std::endl;
    std::cout << "totallength:" << totallength << std::endl;
    std::cout << "tileDataNum:" << tileDataNum << std::endl;
    std::cout << "leftpad:" << leftpad << std::endl;
    std::cout << "rightpad:" << rightpad << std::endl;
    std::cout << "toppad:" << toppad << std::endl;
    std::cout << "bottompad:" << bottompad << std::endl;
    std::cout << "highdimcount:" << highdimcount << std::endl;
    std::cout << "outputlength:" << outputlength << std::endl;
    std::cout << "outputlinealign:" << outputlinealign << std::endl;

    context->SetBlockDim(1);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class ReplicationPad2d : public OpDef {
public:
    explicit ReplicationPad2d(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("paddings")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(ReplicationPad2d);
}
